/*
 * Copyright (c) 2005, 2014 vacoor
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 */
package org.vacoor.nothing.common.misc.util;

import com.google.common.base.Function;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.jvnet.winp.WinProcess;
import org.jvnet.winp.WinpException;
import org.vacoor.nothing.common.util.IOUtils;
import org.vacoor.nothing.common.util.Throwables;

import java.io.*;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.nio.charset.Charset;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 系统进程管理
 * 并不是完全可靠的
 *
 * @author vacoor
 */
public abstract class ProcessTree implements Iterable<ProcessTree.OSProcess> {
    // 系统进程
    public interface OSProcess {

        int getPid();

        String getCommandLine();

        void kill() throws InterruptedException;

    }

    public static ProcessTree get() {
        if (Platform.isWindows()) {
            return new WindowsProcessTree();
        } else if (Platform.isLinux()) {
            return new LinuxProcessTree();
        } else if (Platform.isMac() || Platform.isAIX() || Platform.isFreeBSD() || Platform.isSolaris()) {
            return new BasedPsCommandUnixProcessTree();
        }
        throw new UnsupportedOperationException("不支持当前操作系统");
    }

    /**
     * 获取给定 pid 对应的 {@link OSProcess} 对象
     * 如果不存在返回 null
     *
     * @param pid
     * @return
     */
    public abstract OSProcess get(int pid);

    /**
     * 尝试将 {@link Process} 转换为 {@link OSProcess} 对象
     *
     * @param proc
     * @return
     */
    public abstract OSProcess get(Process proc);

    /**
     * 过滤 {@link OSProcess#getCommandLine()} 能匹配正则表达式匹配的 {@link OSProcess}
     *
     * @param pattern 正则模式
     * @return
     */
    public List<OSProcess> query(Pattern pattern) {
        List<OSProcess> procs = Lists.newLinkedList();

        for (Iterator<OSProcess> it = iterator(); it.hasNext(); ) {
            OSProcess proc = it.next();
            if (pattern.matcher(proc.getCommandLine()).matches()) {
                procs.add(proc);
            }
        }
        return procs;
    }

    public List<OSProcess> query(String regex, int flag) {
        return query(Pattern.compile(regex, flag));
    }

    /**
     * 默认忽略大小写
     *
     * @param regex 正则表达式
     * @return 返回符合条件的 OSProcess List
     */
    public List<OSProcess> query(String regex) {
        return query(regex, Pattern.CASE_INSENSITIVE);
    }

    // ------------------

    /**
     * 系统进程
     */
    abstract class AbstractOSProcess implements OSProcess {
        final int pid;

        AbstractOSProcess(final int pid) {
            this.pid = pid;
        }

        @Override
        public int getPid() {
            return pid;
        }
    }


    // ------------ ProcessTree implementation ----------------

    // ------------------------ windows ---------------------

    /**
     * Windows 进程管理实现, 使用 WinP
     */
    static class WindowsProcessTree extends ProcessTree implements Iterable<OSProcess> {

        @Override
        public OSProcess get(int pid) {
            return new WindowsProcess(new WinProcess(pid));
        }

        @Override
        public OSProcess get(Process process) {
            return new WindowsProcess(new WinProcess(process));
        }

        @Override
        public Iterator<OSProcess> iterator() {
            return Iterators.transform(WinProcess.all().iterator(), new Function<WinProcess, OSProcess>() {
                @Override
                public OSProcess apply(WinProcess input) {
                    return new WindowsProcess(input);
                }
            });
        }

        // 包装 WinProcess,
        class WindowsProcess extends AbstractOSProcess {
            private final WinProcess proc;

            WindowsProcess(WinProcess proc) {
                super(proc.getPid());
                this.proc = proc;
            }

            @Override
            public void kill() throws InterruptedException {
                //proc.kill();
                proc.killRecursively();
            }

            /*
            public void killRecursively() throws InterruptedException {
                proc.killRecursively();
            }
            */

            @Override
            public String getCommandLine() {
                /*
                 * TODO WinP 有个 Bug, 如果程序启动时没有参数, 则调用 getCommandLine 会出错 (测试环境 x86-64)
                 * eg: soffice  getCommandLine 应该是 soffice 但是却抛出异常, 但是 soffice -a --> soffice -a 没问题
                 */
                try {
                    return proc.getCommandLine();
                } catch (WinpException ignore) {
                    return "";
                }
            }

            @Override
            public String toString() {
                return "WindowsProcess{ pid: " + getPid() + ", arguments: " + getCommandLine() + " }";
            }
        }
    }

    // --------------------------- linux ---------------------------
    static abstract class UnixProcessTree extends ProcessTree {
        @Override
        public OSProcess get(int pid) {
            return gatherProcesses().get(pid);
        }

        @Override
        public OSProcess get(Process proc) {
            try {
                return get((Integer) JDKUnixProc.PID_FIELD.get(proc));
            } catch (IllegalAccessException e) { // impossible
                IllegalAccessError x = new IllegalAccessError();
                x.initCause(e);
                throw x;
            }
        }

        @Override
        public Iterator<OSProcess> iterator() {
            return gatherProcesses().values().iterator();
        }

        /**
         * 获取所有进程信息
         */
        abstract Map<Integer/*pid*/, OSProcess> gatherProcesses();

        // 基本的 Unix 进程
        abstract class UnixProcess extends AbstractOSProcess {

            UnixProcess(int pid) {
                super(pid);
            }

            // 获取进程文件夹下相关文件
            protected File getProcFile(String relationPath) {
                return new File(new File("/proc/" + getPid()), relationPath);
            }

            @Override
            public void kill() throws InterruptedException {
                try {
                    int pid = getPid();
                    JDKUnixProc.DESTROY_PROCESS.invoke(null, pid);
                } catch (IllegalAccessException e) {
                    // this is impossible
                    IllegalAccessError x = new IllegalAccessError();
                    x.initCause(e);
                    throw x;
                } catch (InvocationTargetException e) {
                    // tunnel serious errors
                    if (e.getTargetException() instanceof Error)
                        throw (Error) e.getTargetException();
                    // otherwise log and let go. I need to see when this happens
                }
            }

            /*
            @Override
            public void killRecursively() throws InterruptedException {
                kill();
            }
            */

            @Override
            public String toString() {
                return "UnixProcess{ pid: " + getPid() + ", arguments= " + getCommandLine() + " }";
            }
        }
    }

    // 基于 /proc
    static abstract class ProcfsUnixProcessTree extends UnixProcessTree {
        // 从 /proc 下获取所有进程信息
        @Override
        protected Map<Integer, OSProcess> gatherProcesses() {
            Map<Integer, OSProcess> processMap = Maps.newHashMap();
            File[] procs = new File("/proc").listFiles(new FileFilter() {
                @Override
                public boolean accept(File pathname) {
                    return pathname.isDirectory();
                }
            });

            for (File proc : procs) {
                try {
                    int pid = Integer.parseInt(proc.getName());
                    processMap.put(pid, createProcess(pid));
                } catch (NumberFormatException ignore) {
                } catch (IOException ignore2) {
                }
            }

            return Collections.unmodifiableMap(processMap);
        }

        abstract OSProcess createProcess(int pid) throws IOException;
    }

    // unix 可以利用 JDK
    private static final class JDKUnixProc {
        private static final Field PID_FIELD;
        private static final Method DESTROY_PROCESS;

        static {
            try {
                Class<?> clazz = Class.forName("java.lang.UNIXProcess");
                PID_FIELD = clazz.getDeclaredField("pid");
                PID_FIELD.setAccessible(true);

                DESTROY_PROCESS = clazz.getDeclaredMethod("destroyProcess", int.class);
                DESTROY_PROCESS.setAccessible(true);
            } catch (ClassNotFoundException e) {
                LinkageError x = new LinkageError();
                x.initCause(e);
                throw x;
            } catch (NoSuchFieldException e) {
                LinkageError x = new LinkageError();
                x.initCause(e);
                throw x;
            } catch (NoSuchMethodException e) {
                LinkageError x = new LinkageError();
                x.initCause(e);
                throw x;
            }
        }
    }

    // ------- based ps command
    static class BasedPsCommandUnixProcessTree extends UnixProcessTree {
        // pid, arguments
        private static final Pattern PS_PATTERN = Pattern.compile("^\\s*([0-9]+)\\s*(.*)$");
        private static final String[] PS_CMD = {"/bin/ps", "-e", "-o", "pid,args"};

        static {
            File file = new File("/bin/ps");
            if (!file.exists() || !file.canExecute()) {
                throw new IllegalStateException("/bin/ps can't find or is not an executable file");
            }
        }

        @Override
        protected Map<Integer/*pid*/, OSProcess> gatherProcesses() {
            Map<Integer, OSProcess> processMap = Maps.newHashMap();
            try {
                // Process process = new ProcessBuilder(command).start();
                Process proc = Runtime.getRuntime().exec(PS_CMD);
                // proc.waitFor();

                BufferedReader reader = new BufferedReader(new InputStreamReader(proc.getInputStream()));
                String line = reader.readLine();    // ignore title
                while ((line = reader.readLine()) != null) {
                    Matcher matcher = PS_PATTERN.matcher(line);
                    if (matcher.matches()) { // 应该始终是true
                        final int pid = Integer.parseInt(matcher.group(1));
                        final String arguments = matcher.group(2);

                        processMap.put(pid, new UnixProcess(pid) {
                            @Override
                            public String getCommandLine() {
                                return arguments;
                            }
                        });
                    }
                }
                IOUtils.close(reader);
            } catch (Exception e) {
                throw Throwables.unchecked(e);
            }
            return Collections.unmodifiableMap(processMap);
        }
    }

    // -------- linux implementation
    static class LinuxProcessTree extends ProcfsUnixProcessTree {

        @Override
        OSProcess createProcess(int pid) throws IOException {
            return new LinuxProcess(pid);
        }

        class LinuxProcess extends UnixProcess {
            private String arguments;

            LinuxProcess(int pid) {
                super(pid);
            }

            @Override
            public synchronized String getCommandLine() {
                if (arguments != null) {
                    return arguments;
                }

                File cmdline = getProcFile("cmdline");
                try {
                    arguments = IOUtils.readText(cmdline, Charset.defaultCharset());
                } catch (IOException e) {
                    arguments = "";
                }
                return arguments;
            }
        }
    }
}
